################################################################################
class rap_mix:
  """ Stores the mixing matrix from an olg file """
  def __init__(self,filename=None,levs=None,**args):
    self.clear()
    if levs is not None: self.v_levs=levs          # pointer to level object
    if filename: self.load_olg(filename,**args)

# ------------------------------------------------------------------------------

  def clear(self):
    self.v_levs=None
    self.v_mixmat={}               # mixing matrix, index is symmetry tuple
    self.v_symmap={}               # store symmetry tuple for each state

# ------------------------------------------------------------------------------

  def set_levs(self,levs):
    """ Set rap_states object. """
    if not isinstance(levs,rap_states):
      msg='levs must be a rap_states instance (or child instance)'
      raise ErrorRAP(msg)
    self.v_levs=levs

# ------------------------------------------------------------------------------

  def levs(self):
    """ Return rap_states object. """
    return self.v_levs

# ------------------------------------------------------------------------------

  def num_levels(self):
    """ Return number of levels. """
    return len(self.v_levs)

# ------------------------------------------------------------------------------

  def loadfile(self,filename,**args):
    """ Load in mixing data from data file. """
    ft=rlf_filetype(filename)
    if ft == 'olg':
      self.load_olg(filename,**args)
    else:
      msg='do not know how to read file: %s'%filename
      raise ErrorRAP(msg)

# ------------------------------------------------------------------------------

  def symmetries(self,states=None,imap=False):
    """
    Return list of symmetries in olg file. 
    
    Each symmetry is denoted by a tuple containing either (2S+1,L,parity)
    for LS-coupling or (2J,parity) for LSJ-coupling.  If states is given,
    only a subset of symmetries are returned corresponding to the terms in
    the list.
    """
    if imap and states is not None:
      states=map(self.v_levs._mapme,states)
    if not states:
      syms=self.v_mixmat.keys()
    else:
      try:
        syms=[]
        for t in states:
          s=self.symmetry(t)
          if s not in syms: syms.append(s)
      except:
        return False
    syms.sort()
    return syms

# ------------------------------------------------------------------------------

  def symmetry(self,idx,imap=False):
    """ Return symmetry for state with index, idx. """
    tidx=self.v_levs._mapme(idx,imap)
    return self.v_symmap[tidx]

# ------------------------------------------------------------------------------

  def entry(self,sym,imap=False):
    """ Return mixing matrix dictionary entry for single symmetry. """
    if sym not in self.v_mixmat.keys(): return False
    out=self.v_mixmat[sym]
    if not imap: return out
    keys=sorted(out.keys())
    keys=map(self.v_levs.lmap_reverse,keys)
    skeys=sorted(keys)
    tout={}
    for key,val in out.iteritems():
      newkey=self.v_levs.lmap_reverse(key)
      tmp=[]
      for i in range(len(keys)):
        tmp.append(out[key][keys.index(skeys[i])])
      tout[newkey]=tmp
    return tout

# ------------------------------------------------------------------------------

  def coeff(self,idx,imap=False):
    """ Return mixing coefficient dictionary for single state. """
    tidx=self.v_levs._mapme(idx,imap)
    s=self.symmetry(tidx)
    if not s: return False
    keys=self.v_levs.symfilter(s)
    coeffs=self.v_mixmat[s][tidx][:]
    out={}
    for k in keys: 
      kidx=k
      if imap: kidx=self.v_levs.lmap_reverse(kidx)
      out[kidx]=coeffs.pop(0)
    return out

# ------------------------------------------------------------------------------

  def matrix(self,symlist=None,imap=True):
    """ 
    Return 2D list containing mixing matrix for symmetries in symlist.
    """
    # get list of terms having symmetries
    tsymlist=copy.deepcopy(symlist)
    if not symlist:
      tsymlist=self.symmetries()
      termlist=self.v_levs.indices()
    elif symlist in self.symmetries():
      tsymlist=[symlist]
      termlist=self.v_levs.symfilter(symlist)
    else:
      termlist=[]
      for sym in symlist: termlist.extend(self.v_levs.symfilter(sym))
    if imap is True:
      termlist=map(self.v_levs.lmap_reverse,termlist)
    termlist.sort()

    # prepare final matrix and dummy row
    trow=[0 for t in termlist]
    tmat=[trow for t in termlist]

    # construct matrix
    for s in tsymlist:
      sterms=sorted(self.v_levs.symfilter(s,imap=imap))
      srow=self.entry(s,imap=imap)
      for tr in sterms:
        ridx=termlist.index(tr)
        row=trow[:]
        for i in range(len(termlist)):
          tc=termlist[i]
          if tc in sterms:
            j=sterms.index(tc)
            row[i]=srow[tr][j]
          else:
            row[i]=0.0
        tmat[ridx]=row
    return tmat

# ------------------------------------------------------------------------------

  def load_olg(self,filename,cup=None):
    """ Read information from olg file """
    # first some useful functions
    def get_sym(line,cup):
      # get symmetry tuple from line (global variable)
      try:
        if cup == 'LS':
          sym=(int(line[4:7]),int(line[7:10]),int(line[10:13]))
        else:
          sym=(int(line[8:11]),int(line[17:20]))
      except:
        return False
      return sym

    if self.levs() == None: self.v_levs=rlf_loadstates(filename,cup)
    tcup=self.v_levs.coupling()
    if cup is not None and cup != tcup:
      msg='given coupling scheme inconsistent with levels object'
      raise ErrorRAP(msg)
    cup=tcup

    # try to open olg file
    if rlf_filetype(filename) != 'olg': return False
    try:
      fin=open(filename,'r')
    except:
      return False

    # find start of mixing matrices
    line=fin.readline()
    while 'MATRIX' not in line: 
      if line == '': return False
      line=fin.readline()
    istart=51
    if cup == 'LSJ':
      istart=55
      line=fin.readline()
      while 'MATRIX' not in line: 
        if line == '': return False
        line=fin.readline()
    line=fin.readline()
    line=fin.readline()

    # read in mixing matrices
    termsym=[]                   # store symmetry for each term
    oldsym=get_sym(line,cup)
    idx=1
    termsym.append((idx,oldsym))
    cstr=line[istart:]
    alldata={}
    tmat=[]
    tdict={}
    done=False
    while not done:
      line=fin.readline()
      if 'CORE' in line: done=True
      sym=get_sym(line,cup)
      if not sym and not done:
        cstr+=line[istart:]
        continue
      cstr=cstr.strip()[:-8]
      coeffs=cstr.split()
      coeffs=map(float,coeffs)
      tdict[idx]=coeffs
      if done or sym != oldsym:
        nc=999999
        for t in tdict.keys(): nc=min(nc,len(tdict[t]))
        for t in tdict.keys(): tdict[t]=tdict[t][:nc]
        alldata[oldsym]=tdict
        tdict={}
      if sym:
        idx+=1
        cstr=line[istart:]
        oldsym=sym
        termsym.append((idx,oldsym))

    # get mapping from energy-order to symmetry-order
    while 'K*CM' not in line:
      if line == '': return False
      line=fin.readline()
    map_es=[]
    while True:
      line=fin.readline()
      tmp=line.split()
      try:
        it=int(tmp[1])
      except:
        break
      map_es.append(it)
    fin.close()

    # get reverse map
    map_se=[]
    for i in range(1,len(map_es)+1): 
      map_se.append(map_es.index(i)+1)

    # get list of symmetries for each term
    ttermsym=[]
    for (i,s) in termsym: ttermsym.append((map_se[i-1],s))
    ttermsym.sort()
    termsym={}
    for (i,s) in ttermsym: termsym[i]=s
    self.v_symmap=termsym

    # convert symmetry-ordered indices in mixing matrix to energy-ordered
    for s in alldata.keys():
      new={}
      oldterms=alldata[s].keys()
      oldterms.sort()
      newterms=[]
      for ot in oldterms: newterms.append(map_se[ot-1])
      newterms_sort=newterms[:]
      newterms_sort.sort()
      tmap=[]
      for nt in newterms: tmap.append(newterms_sort.index(nt)+1)
      for t in oldterms:
        mt=map_se[t-1]
        row=alldata[s][t]
        newrow=[0 for i in range(len(row))]
        for i in range(1,len(row)+1):
          newrow[tmap[i-1]-1]=row[i-1]
        new[mt]=newrow
      alldata[s]=new

    self.v_mixmat=alldata
    return True

